

clearvars -except K


%% data and prior
load('../true_DGP/priorWSigma.mat')
data=readtable('../../Data/data_main_sample.csv');
dd=readtable('../../Data/distance_capitals.csv');

data=data(data.year>=1950&data.year<=2000&isfinite(data.D),:);
ccodes=unique(data.ccode);
years=unique(data.year(data.year>1950));
N=length(ccodes);
T=length(years);
y=zeros(N,T); D=y; M=y; X=repmat(y,1,1,K);
for n=1:N
    dn=data(data.ccode==ccodes(n),:);
    for t=1:T
        if K==0
            if isempty(setdiff(years(t),dn.year))
                if isfinite(dn.y(years(t)==dn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                end
            end
        elseif K==1
            if isempty(setdiff(years(t),dn.year))
                if isfinite(dn.y(years(t)==dn.year))&&isfinite(dn.Y(years(t)-1==dn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                    X(n,t,1)=log(dn.Y(years(t)-1==dn.year));
                end
            end
        elseif K==2
            if isempty(setdiff(years(t),dn.year))&&isempty(setdiff(years(t)-1,dn.year))
                if isfinite(dn.y(years(t)==dn.year))&&isfinite(dn.TinP(years(t)==dn.year))&&...
                        isfinite(dn.Y(years(t)-1==dn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                    X(n,t,1)=dn.TinP(years(t)==dn.year);
                    X(n,t,2)=log(dn.Y(years(t)-1==dn.year));
                end
            end
        elseif K==3
            if isempty(setdiff(years(t),dn.year))&&isempty(setdiff(years(t)-1,dn.year))
                if isfinite(dn.y(years(t)==dn.year))&&isfinite(dn.TinP(years(t)==dn.year))&&...
                        isfinite(dn.Y(years(t)-1==dn.year))&&isfinite(dn.Trade(years(t)-1==dn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                    X(n,t,1)=dn.TinP(years(t)==dn.year);
                    X(n,t,2)=log(dn.Y(years(t)-1==dn.year));
                    X(n,t,3)=dn.Trade(years(t)-1==dn.year);
                end
            end
        elseif K==4
            if isempty(setdiff(years(t),dn.year))&&isempty(setdiff(years(t)-1,dn.year))
                if isfinite(dn.y(years(t)==dn.year))&&isfinite(dn.TinP(years(t)==dn.year))&&...
                        isfinite(dn.Y(years(t)-1==dn.year))&&isfinite(dn.Trade(years(t)-1==dn.year))&&...
                        isfinite(dn.yearsDem(years(t)==dn.year))
                    M(n,t)=1; % not missing
                    y(n,t)=dn.y(years(t)==dn.year);
                    D(n,t)=dn.D(years(t)==dn.year);
                    X(n,t,1)=dn.TinP(years(t)==dn.year);
                    X(n,t,2)=log(dn.Y(years(t)-1==dn.year));
                    X(n,t,3)=dn.Trade(years(t)-1==dn.year);
                    X(n,t,4)=dn.yearsDem(years(t)==dn.year);
                end
            end
        end
    end
end
for k=1:K
    X(:,:,k)=M.*(X(:,:,k)-sum(sum(X(:,:,k)))/sum(sum(M==1)));
end
clear dn t k

Z=10^(-6)*table2array(dd(:,2:end));
KZ=size(Z,3);

ep_sd=sqrt(diag(Sigma));

% Sparse-grid integration
[sgi_nodes,sgi_weights]=nwspgr('KPN',2,8);

cnames=cell(N,1);
for n=1:N
    a=unique(data.idacr(data.ccode==ccodes(n)));
    cnames{n}=a{1};
end

ccodes=table(ccodes,cnames);

clear a cnames data dd n


%% for Knitro

% Jacobian sparsity pattern:
JPattern=[repmat(eye(N),T,1),ones(N*T,2),ones(N*T,2*N),ones(N*T,N),...
    repmat(eye(N),T,1),zeros(N*T,N),ones(N*T,K),ones(N*T,KZ),eye(N*T)];
JPattern=sparse(JPattern);

% Hessian sparsity pattern
% log-posterior
HPattern=zeros(6*N+2+K+KZ+N*T,6*N+2+K+KZ+N*T);
  % alpha,theta,beta,v,f
HPattern(1:(5*N+2),1:(5*N+2))=eye(5*N+2);
  % vs
HPattern(5*N+2+(1:N),:)=[zeros(N,5*N+2),eye(N),zeros(N,K+KZ),repmat(eye(N),1,T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=[zeros(N*T,5*N+2),repmat(eye(N),T,1),zeros(N*T,K+KZ),eye(N*T)];
% equilibrium constraints
  % alpha
  HPattern(1:N,:)=HPattern(1:N,:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % theta
  HPattern(N+(1:2),:)=HPattern(N+1,:)+ones(2,6*N+2+K+KZ+N*T);
  % f
  HPattern(4*N+(1:N),:)=HPattern(4*N+(1:N),:)+...
        [repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ),repmat(eye(N),1,T)];
  % beta,v,xi,gamma
for n=2:4
    HPattern((n-1)*N+2+(1:N),:)=HPattern((n-1)*N+2+(1:N),:)+...
        [ones(N,5*N+2),zeros(N),ones(N,K+KZ+N*T)];
end
HPattern(6*N+2+(1:K),:)=HPattern(6*N+2+(1:K),:)+...
        [ones(K,5*N+2),zeros(K,N),ones(K,K+KZ+N*T)];
HPattern(6*N+2+K+(1:KZ),:)=HPattern(6*N+2+K+(1:KZ),:)+...
        [ones(KZ,5*N+2),zeros(KZ,N),ones(KZ,K+KZ+N*T)];
  % kappa
HPattern((6*N+2+K+KZ+1):end,:)=HPattern((6*N+2+K+KZ+1):end,:)+...
    [repmat([repmat(eye(N),1,2),ones(N,2*N+2),eye(N),zeros(N),ones(N,K+KZ)],T,1),eye(N*T)];
HPattern=1*(HPattern>0);

HPattern=sparse(HPattern);

clear n

% options
options=optimset('Display','off','JacobPattern',JPattern,'HessPattern',HPattern);


%% starting values

% "short" replication
load(['K',num2str(K),'nl_replication.mat'],'MAP')
pars0=MAP;

% % "long" replication
% load(['K',num2str(K),'nl_replication.mat'],'srng')
% rng(srng)
% alpha=normrnd(a0,om_a,N,1);
% theta=normrnd(th0,om_th,2,1);
% beta=[normrnd(b0A,om_b,N,1);normrnd(b0D,om_b,N,1)];
% v=gamrnd(s_v,1/d_v,N,1).^(-1);
% f=normrnd(f0,om_f,N,1);
% vs=gamrnd(s_vs,1/d_vs,N,1).^(-1);
% xi=2*(-1).^(rand(K,1)<0.5).*rand(K,1);
% gamma=2*rand(KZ,1);
% kppa=zeros(N*T,1);
% pars0=[alpha;theta;beta;v;f;vs;xi;gamma;kppa];
% clear alpha theta beta v f vs xi gamma kppa 


%% maximum-a-posteriori estimator
[MAP,lp,exit]=knitromatlab(@(x)log_posterior(x,D,M,X,Z,...
    a0,om_a,th0,om_th,b0A,b0D,om_b,s_v,d_v,f0,om_f,s_vs,d_vs),pars0,[],[],...
  [],[],[-Inf(3*N+2,1);zeros(N,1);-Inf(N,1);zeros(N,1);-Inf(K,1);zeros(KZ,1);-Inf(N*T,1)],...
  Inf(length(pars0),1),@(x)eq_con_jac_NOlearn(x,y,D,M,X,Z,sgi_nodes,sgi_weights,ep_sd,Sigma),[],...
  options,'knitro.opt');


%%
clearvars -except K MAP exit
save(['MAPK',num2str(K),'nl.mat'])



